from matplotlib import pyplot as plt
import pandas as pd
import numpy as np
from grey_model import GM11

plt.rcParams['font.sans-serif'] = ['SimHei']

cols = ['year', 'area', 'intensity', 'SOC', 'SIC', 'STC', 'N', 'C/N']
features = ['SOC', 'SIC', 'STC', 'N', 'C/N']
area = ['G17', 'G19', 'G21', 'G6', 'G12', 'G18', 'G8', 'G11', 'G16', 'G9', 'G13', 'G20']
coff_S = {
    'SOC': {'LGI': 0.02, 'MGI': 0.04, "HGI": 0.08, "NG": -0.06},
    'N': {'LGI': 0.01, 'MGI': 0.05, "HGI": -0.09, "NG": 0.0},
    'SIC': {'LGI': 0.0, 'MGI': 0.0, "HGI": 0.08, "NG": 0.0},
    'C/N': {'LGI': 0.0, 'MGI': 0.0, "HGI": 0.08, "NG": 0.0},
    'STC': {'LGI': 0.0, 'MGI': 0.0, "HGI": 0.08, "NG": 0.0},
}


def visualization(a='G17'):
    df = pd.read_excel('data/attach14.xlsx', names=cols)
    df = df[df['area'] == a]
    s = df['intensity'].iloc[0]
    print(df.groupby('year').mean())
    df = df.groupby('year').mean()

    fig, ax = plt.subplots(len(features), 1, sharex=True)
    ax = ax.flatten()
    for i, col in enumerate(features):
        if i == 0:
            plt.title('{}:{}'.format(a, s))
        ax[i].plot(df[col], label=col, color='green', marker='D', linestyle='-.')
        ax[i].legend()
    plt.show()


def grey_predict():
    df = pd.read_excel('data/attach14.xlsx', names=cols)
    result = pd.DataFrame(columns=['area', 'SOC', 'SIC', 'STC', 'N', 'C/N'])
    count = 0
    loss = []
    for area_name in area:
        df_area = df[df['area'] == area_name]
        s = df_area['intensity'].iloc[0]
        df_area = df_area.groupby('year').mean()
        tmp = []

        fig, ax = plt.subplots(len(features), 1, sharex=True)
        ax = ax.flatten()
        for i, col in enumerate(features):
            data = pd.DataFrame(data=df_area[col])
            a = GM11(data, predstep=1, S=coff_S[col][s])
            a.fit()
            # print('GM(1,1)的拟合值是： ', a.fit())
            prediction = a.predict()
            print('{:^5}{:^8}的预测值是: {:>8.4f}, 相对误差校验: {:<.4f}'.format(area_name, col, prediction[0], a.loss()))
            # print('GM(1,1)的预测误差是： ', a.loss())
            # print('errors:', a.errors())
            tmp.extend(prediction)
            count += 1
            loss.append(a.loss())
            raw_data = df_area[col].values.tolist()
            predictions = raw_data + [prediction[0]]
            ax[i].plot(predictions, color='red', label=f'{area_name}小区{col}预测')
            ax[i].plot(raw_data, color='green', label=f'{area_name}小区{col}历史值')
            ax[i].legend(loc=2)
            ax[i].grid()
        # plt.show()
        result = result.append(pd.DataFrame(columns=result.columns, data=[[area_name] + tmp]), ignore_index=True)
    result.to_csv('data/q3_result.csv')
    print('count={}, average of loss={}'.format(count, np.mean(loss)))


def main():
    grey_predict()
    # for a in area:
    #     visualization(a)


if __name__ == "__main__":
    # visualization()
    # add()
    main()
